#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
数据集处理程序
使用 -p 参数指定数据集目录
支持NV12格式.raw文件转换为PNG，以及图片尺寸填充
"""

import argparse
import os
import sys
import cv2
import numpy as np
from pathlib import Path


def validate_directory(directory_path):
    """验证目录是否存在且可访问"""
    if not os.path.exists(directory_path):
        print(f"错误: 目录 '{directory_path}' 不存在")
        return False
    
    if not os.path.isdir(directory_path):
        print(f"错误: '{directory_path}' 不是一个目录")
        return False
    
    if not os.access(directory_path, os.R_OK):
        print(f"错误: 没有权限访问目录 '{directory_path}'")
        return False
    
    return True


def convert_nv12_to_640x640_png(raw_file_path, output_dir, rotation_angle=0):
    """将NV12格式的.raw文件直接转换为640x640的PNG格式，不使用中间文件"""
    try:
        # 读取.raw文件
        with open(raw_file_path, 'rb') as f:
            raw_data = f.read()
        
        # NV12格式：Y分量 + UV分量交错
        # 假设输入是480x480，则Y分量大小为480*480，UV分量大小为480*480//2
        height, width = 480, 480
        y_size = height * width
        uv_size = y_size // 2
        
        if len(raw_data) < y_size + uv_size:
            print(f"警告: {raw_file_path} 文件大小不足，跳过处理")
            return False
        
        # 创建NV12格式的numpy数组
        nv12 = np.frombuffer(raw_data, dtype=np.uint8)
        
        # 重塑为正确的尺寸
        y_plane = nv12[:y_size].reshape(height, width)
        uv_plane = nv12[y_size:y_size+uv_size].reshape(height//2, width)
        
        # 将Y和UV平面合并为OpenCV期望的NV12格式
        nv12_array = np.concatenate([y_plane.flatten(), uv_plane.flatten()])
        nv12_reshaped = nv12_array.reshape(height + height//2, width)
        
        # 转换为BGR（480x480）
        bgr_480 = cv2.cvtColor(nv12_reshaped, cv2.COLOR_YUV2BGR_NV12)
        
        # 如果指定了旋转角度，先进行旋转
        if rotation_angle != 0:
            # 计算旋转中心点
            center = (width // 2, height // 2)
            # 获取旋转矩阵
            rotation_matrix = cv2.getRotationMatrix2D(center, rotation_angle, 1.0)
            # 执行旋转
            bgr_480 = cv2.warpAffine(bgr_480, rotation_matrix, (width, height))
        
        # 创建640x640的黑色画布
        canvas_640 = np.zeros((640, 640, 3), dtype=np.uint8)
        
        # 将480x480的图像直接放置在左上角
        canvas_640[:height, :width] = bgr_480
        
        # 生成输出文件名
        base_name = os.path.splitext(os.path.basename(raw_file_path))[0]
        if rotation_angle != 0:
            output_path = os.path.join(output_dir, f"{base_name}_rot{rotation_angle}_640x640.png")
        else:
            output_path = os.path.join(output_dir, f"{base_name}_640x640.png")
        
        # 保存最终的640x640 PNG文件
        cv2.imwrite(output_path, canvas_640)
        print(f"  转换完成: {os.path.basename(raw_file_path)} -> {os.path.basename(output_path)} (640x640)")
        return True
        
    except Exception as e:
        print(f"  转换失败: {raw_file_path} - {e}")
        return False


def pad_image_to_640x640(image_path, output_dir, rotation_angle=0):
    """将480x480的PNG图片填充到640x640，原图放在左上角，其余部分填充黑色"""
    try:
        # 读取图片
        img = cv2.imread(image_path)
        if img is None:
            print(f"  读取失败: {image_path}")
            return False
        
        # 检查图片尺寸
        h, w = img.shape[:2]
        if h != 480 or w != 480:
            print(f"  跳过: {image_path} 尺寸为 {w}x{h}，不是480x480")
            return False
        
        # 如果指定了旋转角度，先进行旋转
        if rotation_angle != 0:
            # 计算旋转中心点
            center = (w // 2, h // 2)
            # 获取旋转矩阵
            rotation_matrix = cv2.getRotationMatrix2D(center, rotation_angle, 1.0)
            # 执行旋转
            img = cv2.warpAffine(img, rotation_matrix, (w, h))
            print(f"  已旋转 {rotation_angle} 度")
        
        # 创建640x640的黑色画布
        canvas = np.zeros((640, 640, 3), dtype=np.uint8)
        
        # 将原图放在左上角
        canvas[:480, :480] = img
        
        # 生成输出文件名
        base_name = os.path.splitext(os.path.basename(image_path))[0]
        if rotation_angle != 0:
            output_path = os.path.join(output_dir, f"{base_name}_rot{rotation_angle}_640x640.png")
        else:
            output_path = os.path.join(output_dir, f"{base_name}_640x640.png")
        
        # 保存填充后的图片
        cv2.imwrite(output_path, canvas)
        print(f"  填充完成: {os.path.basename(image_path)} -> {os.path.basename(output_path)}")
        return True
        
    except Exception as e:
        print(f"  填充失败: {image_path} - {e}")
        return False


def process_dataset(source_dir, dest_dir, rotation_angle=0):
    """处理数据集目录，转换.raw文件并填充PNG图片"""
    print(f"源目录: {source_dir}")
    print(f"目标目录: {dest_dir}")
    if rotation_angle != 0:
        print(f"旋转角度: {rotation_angle} 度")
    
    # 创建输出目录
    os.makedirs(dest_dir, exist_ok=True)
    print(f"输出目录已创建: {dest_dir}")
    
    # 获取源目录中的所有文件
    try:
        files = os.listdir(source_dir)
        raw_files = [f for f in files if f.endswith('.raw')]
        png_files = [f for f in files if f.endswith('.png')]
        
        print(f"发现 {len(files)} 个文件:")
        print(f"  .raw文件: {len(raw_files)} 个")
        print(f"  .png文件: {len(png_files)} 个")
        
        # 处理.raw文件（直接转换为640x640）
        if raw_files:
            print(f"\n开始转换NV12格式的.raw文件到640x640...")
            converted_count = 0
            for raw_file in raw_files:
                raw_path = os.path.join(source_dir, raw_file)
                if convert_nv12_to_640x640_png(raw_path, dest_dir, rotation_angle):
                    converted_count += 1
            print(f"转换完成: {converted_count}/{len(raw_files)} 个文件")
        
        # 处理PNG文件（如果存在的话）
        if png_files:
            print(f"\n开始填充PNG图片到640x640...")
            padded_count = 0
            for png_file in png_files:
                png_path = os.path.join(source_dir, png_file)
                if pad_image_to_640x640(png_path, dest_dir, rotation_angle):
                    padded_count += 1
            print(f"填充完成: {padded_count}/{len(png_files)} 个文件")
        
        if not raw_files and not png_files:
            print("  未发现.raw或.png文件")
                
    except PermissionError:
        print("错误: 没有权限列出目录内容")
    except Exception as e:
        print(f"错误: 无法读取目录内容: {e}")


def main():
    """主函数"""
    # 创建参数解析器
    parser = argparse.ArgumentParser(
        description="数据集处理程序 - 支持NV12格式转换和图片填充",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
使用示例:
  python yolo5-pre.py -s /path/to/source -d /path/to/destination
  python yolo5-pre.py --source /path/to/source --dest /path/to/destination
  python yolo5-pre.py -s /path/to/source -d /path/to/destination -r 90

功能说明:
  - 自动检测并转换NV12格式的.raw文件直接为640x640的PNG格式
  - 将480x480的PNG图片填充到640x640（原图左上角，其余部分黑色）
  - 支持在填充前对图片进行顺时针旋转（-r参数）
  - 零拷贝处理，直接在内存中完成转换和填充
  - 处理后的文件保存在指定的目标目录中
        """
    )
    
    # 添加必需参数
    parser.add_argument(
        '-s', '--source',
        required=True,
        type=str,
        help='源目录路径（包含.raw和.png文件的目录）'
    )
    
    parser.add_argument(
        '-d', '--dest',
        required=True,
        type=str,
        help='目标目录路径（转换后的文件保存位置）'
    )
    
    # 添加可选参数
    parser.add_argument(
        '-r', '--rotation',
        type=int,
        default=0,
        help='顺时针旋转角度（默认0度，表示不旋转）'
    )
    
    # 解析命令行参数
    args = parser.parse_args()
    
    # 获取源目录和目标目录路径
    source_path = args.source
    dest_path = args.dest
    rotation_angle = args.rotation
    
    # 验证源目录
    if not validate_directory(source_path):
        print(f"错误: 源目录 '{source_path}' 无效")
        sys.exit(1)
    
    # 转换为绝对路径
    source_path = os.path.abspath(source_path)
    dest_path = os.path.abspath(dest_path)
    
    # 处理数据集
    process_dataset(source_path, dest_path, rotation_angle)
    
    print("处理完成!")


if __name__ == "__main__":
    main()
